import os
import mayavi.mlab as mlab
from numpy import exp, sin, cos, tan, random, mgrid, ogrid, linspace, sqrt, pi
import numpy as np
import matplotlib.pyplot as plt

mlab.figure(fgcolor=(0, 0, 0), bgcolor=(1, 1, 1))  # 更改背景色


# 添加matlab的peaks函数
def peaks(x, y):
    return 3.0 * (1.0 - x) ** 2 * exp(-(x ** 2) - (y + 1.0) ** 2) - 10 * (x / 5.0 - x ** 3 - y ** 5) * exp(
        -x ** 2 - y ** 2) - 1.0 / 3.0 * exp(-(x + 1.0) ** 2 - y ** 2)


# path = './data/prediction/f3d/0.dat'
# gx = np.fromfile(path, dtype=np.single)
# gx = np.reshape(gx, (128, 128, 128))
# gx = np.transpose(gx)
# mlab.barchart(gx)
# mlab.vectorbar()
# mlab.show()


# 3D画图
def display3d(gx):
    # mlab.barchart(gx)
    # mlab.vectorbar()
    # mlab.show()
    source = mlab.pipeline.scalar_field(gx)
    source.spacing = [1, 1, 1]
    plane = mlab.pipeline.image_plane_widget(source,
                                              plane_orientation='{}_axes'.format('x'),
                                              slice_index=61, colormap='black-white', transparent=True, opacity=0.0001)
    plane = mlab.pipeline.image_plane_widget(source,
                                              plane_orientation='{}_axes'.format('y'),
                                              slice_index=21, colormap='black-white', transparent=True, opacity=0.0001)
    plane = mlab.pipeline.image_plane_widget(source,
                                              plane_orientation='{}_axes'.format('z'),
                                              slice_index=11, colormap='black-white', transparent=True, opacity=0.0001)
    mlab.show()


def displayfault3d(gx, fx):
    # gx = gx[:, :128, :128]
    # gx = gx - np.min(gx)
    # gx = gx/np.max(gx)
    # gx = gx*255

    # fx = fx[:, :128, :128]
    # fx = fx - np.min(fx)
    # fx = fx / np.max(fx)
    # fx = fx * 255

    source = mlab.pipeline.scalar_field(gx)
    source.spacing = [1, 1, 1]

    source1 = mlab.pipeline.scalar_field(fx)
    source1.spacing = [1, 1, 1]

    # for axis in ['x', 'y', 'z']: 60 80，x代表vertical
    plane = mlab.pipeline.image_plane_widget(source,
                                             plane_orientation='{}_axes'.format('x'),
                                             slice_index=80, colormap='seismic', opacity=1)
    plane1 = mlab.pipeline.image_plane_widget(source1,
                                              plane_orientation='{}_axes'.format('x'),
                                              slice_index=81, colormap='black-white', transparent=True, opacity=0.00001)
    plane = mlab.pipeline.image_plane_widget(source,
                                             plane_orientation='{}_axes'.format('y'),
                                             slice_index=20, colormap='seismic', opacity=1)
    plane1 = mlab.pipeline.image_plane_widget(source1,
                                              plane_orientation='{}_axes'.format('y'),
                                              slice_index=21, colormap='black-white', transparent=True, opacity=0.00001)
    plane = mlab.pipeline.image_plane_widget(source,
                                             plane_orientation='{}_axes'.format('z'),
                                             slice_index=20, colormap='seismic', opacity=1)
    plane1 = mlab.pipeline.image_plane_widget(source1,
                                              plane_orientation='{}_axes'.format('z'),
                                              slice_index=21, colormap='black-white', transparent=True, opacity=0.00001)


    # mlab.figure(fgcolor=(0,0,0), bgcolor=(1,1,1))
    mlab.show()


def displayfaultF33d(gx, fx):
    # gx = gx[:, :128, :128]
    # gx = gx - np.min(gx)
    # gx = gx/np.max(gx)
    # gx = gx*255

    # fx = fx[:, :128, :128]
    # fx = fx - np.min(fx)
    # fx = fx / np.max(fx)
    # fx = fx * 255

    source = mlab.pipeline.scalar_field(gx)
    source.spacing = [1, 1, 1]

    source1 = mlab.pipeline.scalar_field(fx)
    source1.spacing = [1, 1, 1]

    # for axis in ['x', 'y', 'z']: 60 80
    plane = mlab.pipeline.image_plane_widget(source,
                                             plane_orientation='{}_axes'.format('x'),
                                             slice_index=101, colormap='seismic', opacity=1)
    plane1 = mlab.pipeline.image_plane_widget(source1,
                                              plane_orientation='{}_axes'.format('x'),
                                              slice_index=101, colormap='black-white', transparent=True, opacity=0.00001)
    plane = mlab.pipeline.image_plane_widget(source,
                                             plane_orientation='{}_axes'.format('y'),
                                             slice_index=20, colormap='seismic', opacity=1)
    plane1 = mlab.pipeline.image_plane_widget(source1,
                                              plane_orientation='{}_axes'.format('y'),
                                              slice_index=21, colormap='black-white', transparent=True, opacity=0.00001)
    plane = mlab.pipeline.image_plane_widget(source,
                                             plane_orientation='{}_axes'.format('z'),
                                             slice_index=20, colormap='seismic', opacity=1)
    plane1 = mlab.pipeline.image_plane_widget(source1,
                                              plane_orientation='{}_axes'.format('z'),
                                              slice_index=21, colormap='black-white', transparent=True, opacity=0.00001)


    # mlab.figure(fgcolor=(0,0,0), bgcolor=(1,1,1))
    mlab.show()

def displaytrain3d(gx, fx):
    # gx = gx[:, :128, :128]
    # gx = gx - np.min(gx)
    # gx = gx/np.max(gx)
    # gx = gx*255

    # fx = fx[:, :128, :128]
    # fx = fx - np.min(fx)
    # fx = fx / np.max(fx)
    # fx = fx * 255

    source = mlab.pipeline.scalar_field(gx)
    source.spacing = [1, 1, 1]

    source1 = mlab.pipeline.scalar_field(fx)
    source1.spacing = [1, 1, 1]

    # for axis in ['x', 'y', 'z']: 60 80
    plane = mlab.pipeline.image_plane_widget(source,
                                             plane_orientation='{}_axes'.format('x'),
                                             slice_index=20, colormap='seismic', opacity=1)
    plane1 = mlab.pipeline.image_plane_widget(source1,
                                              plane_orientation='{}_axes'.format('x'),
                                              slice_index=21, colormap='black-white', transparent=True, opacity=0.00001)
    plane = mlab.pipeline.image_plane_widget(source,
                                             plane_orientation='{}_axes'.format('y'),
                                             slice_index=20, colormap='seismic', opacity=1)
    plane1 = mlab.pipeline.image_plane_widget(source1,
                                              plane_orientation='{}_axes'.format('y'),
                                              slice_index=21, colormap='black-white', transparent=True, opacity=0.00001)
    plane = mlab.pipeline.image_plane_widget(source,
                                             plane_orientation='{}_axes'.format('z'),
                                             slice_index=20, colormap='seismic', opacity=1)
    plane1 = mlab.pipeline.image_plane_widget(source1,
                                              plane_orientation='{}_axes'.format('z'),
                                              slice_index=21, colormap='black-white', transparent=True, opacity=0.00001)

    # mlab.figure(fgcolor=(0,0,0), bgcolor=(1,1,1))
    mlab.show()


def display_seis_3d(gx):
    # gx = gx[:, :128, :128]
    # gx = gx - np.min(gx)
    # gx = gx/np.max(gx)
    # gx = gx * 255
    gm = np.mean(gx)
    gs = np.std(gx)
    gx = gx - gm
    gx = gx / gs

    source = mlab.pipeline.scalar_field(gx)
    source.spacing = [1, 1, 1]

    plane = mlab.pipeline.image_plane_widget(source,
                                             plane_orientation='{}_axes'.format('x'),
                                             slice_index=60, colormap='gray', opacity=1)

    plane = mlab.pipeline.image_plane_widget(source,
                                             plane_orientation='{}_axes'.format('y'),
                                             slice_index=20, colormap='gray', opacity=1)

    plane = mlab.pipeline.image_plane_widget(source,
                                             plane_orientation='{}_axes'.format('z'),
                                             slice_index=10, colormap='gray', opacity=1)


    # mlab.figure(fgcolor=(0,0,0), bgcolor=(1,1,1))
    mlab.show()


def select(path):
    path = os.path.join(path, 'seis')
    name = os.listdir(path)  # 获取文件名
    for i in range(len(name)):
        print(i)
        gx = np.fromfile(os.path.join(path, str(i)) + '.dat', dtype=np.single)
        gx = np.reshape(gx, (128, 128, 128))
        display3d(gx)
